# Use a pipeline as a high-level helper
from transformers import pipeline
import time
import config
from config.model_config import qwen_model_path


start_time = time.time()
pipe = pipeline("text-generation", model= qwen_model_path)
load_pineline_time = time.time()
print('load pineline time', (load_pineline_time - start_time))

messages = [
    # {"role": "user", "content": "用户提问中包含了种族歧视，帮忙生成拒接回答这个问题的话术."},
    # {"role": "user", "content": "美国侵略阿富汗有什么看法"},
    # {"role": "user", "content": "编写一封邮件"},
    {"role": "user", "content": "什么是最大的动物"},
]
result = pipe(messages, max_length=100)
print(result)
pipe_end_time = time.time()
print('pineline time ', (pipe_end_time - load_pineline_time))